import xml.etree.cElementTree as et
import os
import cv2
import utils
import math
import numpy as np
import random
from kmeans import kmeans, avg_iou

classes = [
    "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat",
    "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person",
    "pottedplant", "sheep", "sofa", "train", "tvmonitor"
]

# usage:
# 先使用 xml_file_pars生成标签
# test_datafile 测试数据是否制作正确
# get_train_test_dataset 划分测试集和训练集
# 制作建议框


def xml_file_pars(rootpath, saroot='./tag/'):
    # 遍历文件夹
    utils.ensure_folder(rootpath)
    utils.ensure_folder(saroot)
    anno_path = fr"{rootpath}/Annotations"

    fp = open(fr"{saroot}/alldata.txt", 'w')

    data_list = []

    for xml_file in os.listdir(anno_path):
        data_list = []
        xml_file_path = fr"{anno_path}\{xml_file}"
        # 添加地址
        try:
            tree = et.parse(xml_file_path)
            root = tree.getroot()
            # print(root.find('size').text)
            img_info = root.find('size')
            img_width = int(img_info.find('width').text)
            img_height = int(img_info.find('height').text)
            img_depth = int(img_info.find('depth').text)

            if img_depth != 3:
                print(fr"{xml_file_path} : error!channel does not match")
                continue
            jpg_name = root.find("filename").text
            data_list.append(fr"{rootpath}\JPEGImages\{jpg_name}")
        except Exception as ep:
            print(f"{xml_file}:error! {ep}")
            continue
        # 获取最大边长
        max_side = max(img_width, img_height)
        # 填充为正方形，则要重新更改原点坐标
        new_startx = (max_side - img_width) // 2
        new_starty = (max_side - img_height) // 2
        # 获取缩放比例
        scale = 416 / max_side

        for object in root.findall('object'):
            cls = object.find('name').text
            try:
                bndbox = object.find("bndbox")
                pts0_x = int(bndbox.find('xmin').text)
                pts0_y = int(bndbox.find('ymin').text)
                pts1_x = int(bndbox.find('xmax').text)
                pts1_y = int(bndbox.find('ymax').text)
            except Exception as ep:
                print(fr"{xml_file_path}:wrong! {ep}")
                print(pts0_x, pts0_y, pts1_x, pts1_y)
            w, h = pts1_x - pts0_x, pts1_y - pts0_y
            cent_x, cent_y = (pts1_x + pts0_x) / 2, (pts1_y + pts0_y) / 2

            data_list.append(classes.index(cls))  # 添加类别
            data_list.append(math.ceil((cent_x + new_startx) * scale))  # 添加中心点
            data_list.append(math.ceil((cent_y + new_starty) * scale))
            data_list.append(math.ceil(w * scale))  # 添加宽高
            data_list.append(math.ceil(h * scale))
        print(data_list, cls, classes[data_list[1]])
        for data in data_list:
            fp.write(str(data) + ' ')
        fp.write('\n')
        fp.flush()


def test_datafile(filepath):
    data = []
    with open(filepath, 'r') as fp:
        data = fp.readlines()

    print(data[0].split())
    for item in data:
        item = item.split()
        img_path = item[0]

        boxes = np.array([float(n) for n in item[1:]])
        boxes = np.split(boxes, len(boxes) // 5)

        img = cv2.imread(img_path)
        h, w = img.shape[0], img.shape[1]
        max_side = max(h, w)
        new_img = np.zeros((max_side, max_side, 3), dtype=np.uint8)

        start_y = (max_side - h) // 2
        start_x = (max_side - w) // 2
        new_img[start_y:start_y + h, start_x:start_x + w] += img.copy()
        new_img = cv2.resize(new_img, (416, 416))
        print(img_path)
        print(item)
        for box in boxes:
            obj_cls = int(box[0])
            obj_x = int(box[1])
            obj_y = int(box[2])
            obj_w = int(box[3])
            obj_h = int(box[4])
            # print(img_path)

            # print(img.shape)

            # cv2.imshow("src", img)

            rec_img = cv2.rectangle(new_img, (math.ceil(obj_x - obj_w / 2), math.ceil(obj_y - obj_h / 2)),
                                    (math.ceil(obj_x + obj_w / 2), math.ceil(obj_y + obj_h / 2)), (0, 0, 255), 1)
            rec_img = cv2.putText(rec_img, f'{classes[obj_cls]}', (obj_x, obj_y), cv2.FONT_HERSHEY_COMPLEX, 1,
                                  (0, 255, 0), 2)
        cv2.imshow("new", rec_img)
        cv2.waitKey(0)


# 划分测试集，训练集 比例为2：8
def get_train_test_dataset(srcpath, dstpath='./voc2012'):
    with open(srcpath, 'r') as fp:
        data = fp.readlines()

    utils.ensure_folder(dstpath)
    random.shuffle(data)
    datalen = len(data)
    train_data_len = int(datalen * 0.8)

    train_data = data[0:train_data_len]
    test_data = data[train_data_len:]

    with open(fr"{dstpath}/traindata.txt", 'w') as fp:
        for item in train_data:
            fp.write(item)
    with open(fr"{dstpath}/testdata.txt", 'w') as fp:
        for item in test_data:
            fp.write(item)

    return


# 获取建议框
def get_suggesion_box(filepath, need_kmeans=False,sa_path='./voc2012'):
    with open(filepath, 'r') as fp:
        data = fp.readlines()
    tag_box = []
    out = [[165., 243.]
        , [55., 119.]
        , [100., 174.]
        , [90., 70.]
        , [228., 129.]
        , [15., 22.]
        , [45., 31.]
        , [312., 275.]
        , [30., 62.]]
    for item in data:
        item = item.split()
        boxes = np.array([float(n) for n in item[1:]])
        boxes = np.split(boxes, len(boxes) // 5)
        for box in boxes:
            obj_w = int(box[3])
            obj_h = int(box[4])
            tag_box.append([obj_w / 416, obj_h / 416])
    tag_box = np.array(tag_box)
    print('get tag ok:', tag_box.shape)
    if need_kmeans == True:
        out = kmeans(np.array(tag_box), k=9)

        print(out * 416)
        print("Accuracy: {:.2f}%".format(avg_iou(tag_box, out) * 100))
        out = out * 416
    else:
        out = np.array(out)
        print(out)
        print("Accuracy: {:.2f}%".format(avg_iou(tag_box, out / 416) * 100))



    # 对面积进行排序
    area = out[:,0]*out[:,1]
    sort_index = np.argsort(area)

    #将参数保存
    utils.ensure_folder(sa_path)
    with open(fr"{sa_path}/anchor box.txt",'w') as fp:
        out = out[sort_index].tolist()
        fp.write('[')
        for data in out:
            fp.write(str(data)+',')
        fp.write(']')
        print(f"save {sa_path}/anchor box.txt ok ")
    print(out)
    # print(tag_box)


if __name__ == '__main__':
    # xml_file_pars(r"E:\YOLO\VOCdevkit\VOC2012")
    # test_datafile(r"D:\1.课程记录\Prj_yolo\data\tag\alldata.txt")
    # get_train_test_dataset(r"D:\1.课程记录\Prj_yolo\data\tag\alldata.txt")
    get_suggesion_box(r"D:\1.课程记录\Prj_yolo\data\tag\alldata.txt")
    pass
